{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Introduction to PyTorch\n",
    "=============\n",
    "\n",
    "To familiarize ourselves with [PyTorch](http://pytorch.org/about/) and verify that PyTorch is installed and functioning correctly with CUDA support on Jetson TX1/TX2, first let's run some simple Python commands that load PyTorch and perform some basic operations on tensors.  Then we'll create and run a simple neural network.\n",
    "\n",
    "Loading PyTorch\n",
    "-----------------\n",
    "In order to load PyTorch, first we `import torch` from a Python terminal:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If there's no output from the above command, the PyTorch module was installed and able to be loaded successfully.  \n",
    "If the module is reported to be missing or there are otherwise errors printed, there was a problem building/installing PyTorch.\n",
    "\n",
    "Next let's print out and confirm the version of PyTorch that was installed (the output should read `0.3.0b0+af3964a`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(torch.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Verifying CUDA\n",
    "-----------------\n",
    "Confirm that CUDA + cuDNN support has been installed and is available using the PyTorch `torch.cuda.is_available()` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('CUDA available: ' + str(torch.cuda.is_available()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, allocate some simple tensors and confirm that CUDA arithmatic operations are working:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = torch.cuda.FloatTensor(2).zero_()\n",
    "print('Tensor a = ' + str(a))\n",
    "\n",
    "b = torch.randn(2).cuda()\n",
    "print('Tensor b = ' + str(b))\n",
    "\n",
    "c = a + b\n",
    "print('Tensor c = ' + str(c))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A [tensor](http://pytorch.org/docs/master/tensors.html) is an N-dimensional vector or matrix containing elements of a shared data type (like `float` or `half`)  \n",
    "The `cuda` qualifiers above mean that these tensors are allocated on the GPU, as opposed to CPU.\n",
    "\n",
    "<img src=\"http://pytorch.org/static/img/tensor_illustration.png\">\n",
    "\n",
    "For the future, when we're using CUDA, set the default PyTorch tensor type to CUDA tensors so they are automatically allocated on GPU:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "use_cuda = torch.cuda.is_available()\n",
    "\n",
    "FloatTensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor\n",
    "LongTensor = torch.cuda.LongTensor if use_cuda else torch.LongTensor\n",
    "ByteTensor = torch.cuda.ByteTensor if use_cuda else torch.ByteTensor\n",
    "Tensor = FloatTensor\n",
    "\n",
    "if use_cuda:\n",
    "    torch.set_default_tensor_type('torch.cuda.FloatTensor')\n",
    "    \n",
    "# verify that the default type is CUDA tensor\n",
    "d = FloatTensor([[1, 2, 3], [4, 5, 6]])\n",
    "print('Tensor d = ' + str(d))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that now the `torch.cuda.FloatTensor` type is created without explicitly stating it, even though it was created as generic `FloatTensor`.\n",
    "\n",
    "Next let's create a simple neural network in PyTorch to demonstrate the [`torch.nn`](http://pytorch.org/docs/master/nn.html) package and to test training & inference."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Neural Network Example\n",
    "-----------------"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PyTorch employs a tape-based automatic differention (autograd) system, meaning that it 'records' your dynamic pipeline for backpropagation later on (training) provided a loss function.  The primary package that implements neural networks and associated layer kernels is [`torch.nn`](http://pytorch.org/docs/master/nn.html).  Let's start by importing `torch.nn` along with PyTorch's `Variable` object, which is used to wrap tensors so PyTorch can record the actions performed on them during the pipeline so they can be recalled during the backwards training pass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "from torch.autograd import Variable"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's allocate some tensor variables to hold the network input and outputs.  \n",
    "Note that normally a stored dataset or sensor data would be loaded here, however for simplicity, we are just setting them to be random."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# N is batch size; D_in is input dimension;\n",
    "# H is hidden dimension; D_out is output dimension.\n",
    "N, dim_in, H, dim_out = 64, 1000, 100, 10\n",
    "\n",
    "# Create random Tensors to hold inputs and expected outputs, and wrap them in Variables.\n",
    "x = Variable(torch.randn(N, dim_in))\n",
    "y = Variable(torch.randn(N, dim_out), requires_grad=False)\n",
    "\n",
    "print('x = ' + str(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the tensor type is still `cuda.FloatTensor` (on the GPU).  \n",
    "\n",
    "Now for the fun part, we create our own simple neural network.  This is just using 1D fully-connected layers ([`nn.Linear`](http://pytorch.org/docs/master/nn.html#linear)).  \n",
    "Later on for the DQN, we will utilize the more advanced [`nn.Conv2D`](http://pytorch.org/docs/master/nn.html#conv2d) layers that can process images (deep convolutional networks)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the nn package to define our model as a sequence of layers. nn.Sequential\n",
    "# is a Module which contains other Modules, and applies them in sequence to\n",
    "# produce its output. Each Linear Module computes output from input using a\n",
    "# linear function, and holds internal Variables for its weight and bias.\n",
    "model = nn.Sequential(\n",
    "    nn.Linear(dim_in, H),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(H, dim_out),\n",
    ")\n",
    "\n",
    "# Migrate the model to use CUDA\n",
    "if use_cuda:\n",
    "    model.cuda()\n",
    "    \n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get the neural network model running on the GPU, you just have to call `model.cuda()` on it.  \n",
    "Between that and having the tensors allocated with CUDA as above, that's all we need to do to have the example running on GPU.\n",
    "\n",
    "Next, we define our loss function, which compares the output of the network to the expected output, and informs the backpropagation system of the gradient updates to apply during the training iterations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The nn package also contains definitions of popular loss functions; in this\n",
    "# case we will use Mean Squared Error (MSE) as our loss function.\n",
    "loss_fn = nn.MSELoss(size_average=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally comes the training loop.  This is where we have the current network compute it's prediction of the desired output `y` given the input `x` (the foward pass), then compute the loss versus the expected output before performing backpropagation and gradient descent to decrease the loss and improve the accuracy of the network.  What you should see from the console output, is the loss gradually decreasing over time, meaning our network is learning properly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = 1e-4\n",
    "\n",
    "for t in range(500):\n",
    "    # Forward pass: compute predicted y by passing x to the model. Module objects\n",
    "    # override the __call__ operator so you can call them like functions. When\n",
    "    # doing so you pass a Variable of input data to the Module and it produces\n",
    "    # a Variable of output data.\n",
    "    y_pred = model(x)\n",
    "\n",
    "    # Compute and print loss. We pass Variables containing the predicted and true\n",
    "    # values of y, and the loss function returns a Variable containing the\n",
    "    # loss.\n",
    "    loss = loss_fn(y_pred, y)\n",
    "    print('iteration = %03i, loss = ' % t + str(loss.data[0]))\n",
    "\n",
    "    # Zero the gradients before running the backward pass.\n",
    "    model.zero_grad()\n",
    "\n",
    "    # Backward pass: compute gradient of the loss with respect to all the learnable\n",
    "    # parameters of the model. Internally, the parameters of each Module are stored\n",
    "    # in Variables with requires_grad=True, so this call will compute gradients for\n",
    "    # all learnable parameters in the model.\n",
    "    loss.backward()\n",
    "\n",
    "    # Update the weights using gradient descent. Each parameter is a Variable, so\n",
    "    # we can access its data and gradients like we did before.\n",
    "    for param in model.parameters():\n",
    "        param.data -= learning_rate * param.grad.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can try adjusting the `learning_rate` and number of training passes (set to `500` by default for illustrative purposes) to see how it impacts the training."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
